#!/usr/bin/python3

import json
import socket
import sys

class QMPClient:
    def __init__(self, socket_path):
        self.socket_path = socket_path
        self.sock = None

    def connect(self):
        self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.sock.connect(self.socket_path)
        self._initialize()

    def _initialize(self):
        self._read_response()
        self._send_command({"execute": "qmp_capabilities"})

    def _read_response(self):
        response = b""
        while True:
            chunk = self.sock.recv(4096)
            if not chunk:
                break
            response += chunk
            if b'\n' in chunk:
                break
        return json.loads(response.decode('ascii'))


    def _send_command(self, command):
        self.sock.sendall(json.dumps(command).encode('utf-8') + b'\n')
        return self._read_response()

    def add_virtio_dev_pci(self, memdev, id):
       # Add the virtio-mem-pci device
        device_command = {
               "execute": "device_add",
               "arguments": {
                   "driver": "virtio-mem-pci",
                   "id": id,
                   "memdev": memdev,
                   "bus": "pci.1",
                   }
               }
        return self._send_command(device_command)

    def add_virtio_mem_device(self, size):
        # Add the memory backend
        memdev_command = {
                "execute": "object-add",
                "arguments": {
                    "qom-type": "memory-backend-ram",
                    "id": "virtiomem",
                    "size": size * 1024 * 1024
                    }
        }
        memdev_response = self._send_command(memdev_command)
        print("Memory backend add response:", memdev_response)

        return self.add_virtio_dev_pci("virtiomem", "virtiomem0")

    def set_virtio_mem_size(self, device_id, size):
        command = {
            "execute": "qom-set",
            "arguments": {
                "path": device_id,
                "property": "requested-size",
                "value": size * 1024 * 1024
            }
        }
        return self._send_command(command)

    def get_virtio_mem_size(self, device_id):
        command = {
            "execute":"qom-get",
            "arguments": {
                "path": device_id,
                "property": "requested-size",
            }
        }
        return self._send_command(command)


    def query_virtio(self):
        command = {
                "execute":"query-memory-devices"
                }
        return self._send_command(command)

    def close(self):
        if self.sock:
            self.sock.close()

if __name__ == "__main__":
    qmp_socket_path = "/tmp/qmp-sock"
    device_id = "virtio-mem0"

    client = QMPClient(qmp_socket_path)
    client.connect()
    
    if sys.argv[1] == "qv":
        response = client.query_virtio()
        print(response)
    elif sys.argv[1] == "addm":
        size = int(sys.argv[2])
        response = client.add_virtio_mem_device(size)
        print(response)
    elif sys.argv[1] == "addp":
        response = client.add_virtio_dev_pci("virtiomem", "virtiomem0")
        print(response)

    elif sys.argv[1] == "set":
        # Set memory
        new_size = int(sys.argv[2])
        response = client.set_virtio_mem_size("virtiomem0", new_size)
        print("Set memory response:", response)
    elif sys.argv[1] == "get":
        # Get memory
        response = client.get_virtio_mem_size("virtiomem0")
        print("Get memory response:", response)
    
    client.close()
